﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;

namespace WordpressRestApi
{
    /// <summary>Contains extension methods for <see cref="IAsyncEnumerable{T}"/>.</summary>
    public static class AsyncEnumerable
    {
        #region Inner Types
        private class TransformAsyncEnumerable<T, U> : IAsyncEnumerable<U>
        {
            private readonly IAsyncEnumerable<T> _source;
            private readonly Func<IEnumerable<T>, IEnumerable<U>> _transform;

            public TransformAsyncEnumerable(IAsyncEnumerable<T> source, Func<IEnumerable<T>, IEnumerable<U>> transform)
            {
                if (source == null)
                {
                    throw new ArgumentNullException(nameof(source));
                }
                if (transform == null)
                {
                    throw new ArgumentNullException(nameof(transform));
                }

                _source = source;
                _transform = transform;
            }

            IAsyncEnumerator<U> IAsyncEnumerable<U>.GetAsyncEnumerator()
            {
                return new TransformAsyncEnumerator<T, U>(_source.GetAsyncEnumerator(), _transform);
            }
        }

        private class TransformAsyncEnumerator<T, U> : IAsyncEnumerator<U>
        {
            private readonly IAsyncEnumerator<T> _source;
            private readonly Func<IEnumerable<T>, IEnumerable<U>> _transform;

            public TransformAsyncEnumerator(IAsyncEnumerator<T> source, Func<IEnumerable<T>, IEnumerable<U>> transform)
            {
                if (source == null)
                {
                    throw new ArgumentNullException(nameof(source));
                }
                if (transform == null)
                {
                    throw new ArgumentNullException(nameof(transform));
                }

                _source = source;
                _transform = transform;
            }

            IEnumerable<U> IAsyncEnumerator<U>.CurrentBatch
            {
                get
                {
                    var sourceBatch = _source.CurrentBatch;
                    var transformed = _transform(sourceBatch);

                    return transformed;
                }
            }

            Task<bool> IAsyncEnumerator<U>.MoveNextAsync()
            {
                return _source.MoveNextAsync();
            }
        }
        #endregion

        /// <summary>Enumerates all elements from the enumerable.</summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="enumerable"></param>
        /// <param name="takeCount">Maximum number of elements to retrieve.</param>
        /// <returns>Array of all the elements.</returns>
        public static async Task<T[]> ToArrayAsync<T>(
            this IAsyncEnumerable<T> enumerable,
            int? takeCount = null)
        {
            var enumerator = enumerable.GetAsyncEnumerator();
            var elements = new List<T>();

            while ((takeCount == null || elements.Count < takeCount.Value)
                && await enumerator.MoveNextAsync())
            {
                elements.AddRange(enumerator.CurrentBatch);
            }

            if (takeCount == null)
            {
                return elements.ToArray();
            }
            else
            {
                return elements.Take(takeCount.Value).ToArray();
            }
        }

        /// <summary>Projects an enumeration.</summary>
        /// <typeparam name="T">Source enumeration type.</typeparam>
        /// <typeparam name="U">Target enumeration type.</typeparam>
        /// <param name="enumerable">The enumeration to project.</param>
        /// <param name="projection">The projection of an element.</param>
        /// <returns>Projected enumeration.</returns>
        public static IAsyncEnumerable<U> Select<T, U>(this IAsyncEnumerable<T> enumerable, Func<T, U> projection)
        {
            Func<IEnumerable<T>, IEnumerable<U>> transform = (source) => source.Select(projection);

            return new TransformAsyncEnumerable<T, U>(enumerable, transform);
        }

        /// <summary>Filters an enumeration.</summary>
        /// <typeparam name="T">Source enumeration type.</typeparam>
        /// <param name="enumerable">The enumeration to project.</param>
        /// <param name="predicate"></param>
        /// <returns>Filtered enumeration.</returns>
        public static IAsyncEnumerable<T> Where<T>(this IAsyncEnumerable<T> enumerable, Func<T, bool> predicate)
        {
            Func<IEnumerable<T>, IEnumerable<T>> transform = (source) => source.Where(predicate);

            return new TransformAsyncEnumerable<T, T>(enumerable, transform);
        }
    }
}